import os
import numpy as np
import cv2
import pickle
from transforms3d.euler import mat2euler, euler2mat
from process_utils.process_data import read_pose_from_json, read_object_pcd, read_model_pcd
from process_utils.pose import get_pose_diff
from process_utils.visualization import draw_curves, pcd_bbox_visualization
from optimizer.torch_geometry_optimizer import TorchOptimizerGeometry
import open3d as o3d


def track_ShapeAlign(video_dir, frame_ids, model_points, sym_axis=None):
    object_pointss = []
    pred_poses = []
    gt_poses = []
    energies = []
    t_errs = []
    r_errs = []
    
    for idx in list(frame_ids):
        # read gt pose
        gt_pose_path = os.path.join(video_dir, "objpose", str(idx) + ".json")
        if not os.path.isfile(gt_pose_path):
            gt_pose_path = os.path.join(video_dir, "objpose", str(idx).zfill(5) + ".json")
        gt_pose = read_pose_from_json(gt_pose_path)
        gt_poses.append(gt_pose)

        if idx == frame_ids[0]:  # initialization
            pred_poses.append(gt_pose)
            continue

        # read data
        object_pcd = read_object_pcd(video_dir, idx)
        object_points = np.float32(object_pcd.points)
        object_pointss.append(object_points)

        # predict current pose
        init_state = np.zeros(6)
        init_state[0:3] = pred_poses[-1]["translation"].reshape(3)
        ai, aj, ak = mat2euler(pred_poses[-1]["rotation"])
        init_state[3:6] = np.float32((ai, aj, ak))
        optimizer = TorchOptimizerGeometry(points=object_points, model_points=model_points, init_state=init_state, device="cuda:0")
        optimizer.solve(epoch=100)
        answer_state, answer_energy = optimizer.get_answer()
        print("energy:", answer_energy)
        energies.append(answer_energy)

        # statistics
        pred_pose = {
            "rotation": euler2mat(answer_state[3], answer_state[4], answer_state[5]),
            "translation": answer_state[0:3].reshape(3, 1),
            "scale": pred_poses[-1]["scale"]
        }
        pred_poses.append(pred_pose)
        tdiff, rdiff = get_pose_diff(pred_pose, gt_pose, sym_axis=sym_axis)
        print("error:", tdiff, rdiff)
        t_errs.append(tdiff * 100)
        r_errs.append(rdiff)
    
    # save results
    save_dir = os.path.join("./results", video_dir.replace("/", "_"))
    os.makedirs(save_dir, exist_ok=True)
    pickle.dump(pred_poses, open(os.path.join(save_dir, "pred_poses.pkl"), "wb"))
    pickle.dump(gt_poses, open(os.path.join(save_dir, "gt_poses.pkl"), "wb"))
    draw_curves("./results", video_dir.replace("/", "_"), list(np.arange(1, frame_ids.shape[0])), energies=energies, r_errs=r_errs, t_errs=t_errs)

    # save pcd model
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(model_points)
    o3d.io.write_point_cloud(os.path.join(save_dir, "model.ply"), pcd)

    # visualization
    pcd_bbox_visualization(save_dir=save_dir, model_points=model_points, object_pointss=object_pointss, pred_poses=pred_poses[1:], gt_poses=gt_poses[1:])

    return r_errs, t_errs


if __name__ == "__main__":
    # -----------------------------------------------------------------------------------------------------------
    # kettle
    video_dirs = [
        "/share/datasets/HOI4D_overall/ZY20210800001/H1/C12/N31/S165/s02/T1",
        "/share/datasets/HOI4D_overall/ZY20210800001/H1/C12/N40/S165/s04/T1",
        "/share/datasets/HOI4D_overall/ZY20210800001/H1/C12/N32/S200/s02/T2",
        "/share/datasets/HOI4D_overall/ZY20210800001/H1/C12/N45/S200/s05/T2",
        "/share/datasets/HOI4D_overall/ZY20210800001/H1/C12/N37/S200/s03/T2",
        "/share/datasets/HOI4D_overall/ZY20210800001/H1/C12/N11/S203/s03/T2",
        "/share/datasets/HOI4D_overall/ZY20210800001/H1/C12/N26/S165/s01/T1",
        "/share/datasets/HOI4D_overall/ZY20210800001/H1/C12/N29/S165/s01/T1",
        "/share/datasets/HOI4D_overall/ZY20210800001/H1/C12/N41/S200/s04/T2",
        "/share/datasets/HOI4D_overall/ZY20210800001/H1/C12/N50/S203/s02/T2",
        "/share/datasets/HOI4D_overall/ZY20210800002/H2/C12/N15/S196/s04/T2",
    ]
    frame_ids = np.arange(0, 300)
    category = "rigid/Kettle"
    sym_axis = None
    use_gt_model = True  # use gt model or completed model
    # -----------------------------------------------------------------------------------------------------------
    # # -----------------------------------------------------------------------------------------------------------
    # # toy car
    # video_dirs = [
    #     "/share/datasets/HOI4D_overall/ZY20210800001/H1/C1/N35/S104/s01/T2",
    #     "/share/datasets/HOI4D_overall/ZY20210800001/H1/C1/N26/S104/s05/T2",
    #     "/share/datasets/HOI4D_overall/ZY20210800001/H1/C1/N39/S99/s01/T1",
    #     "/share/datasets/HOI4D_overall/ZY20210800001/H1/C1/N15/S99/s04/T1",
    #     "/share/datasets/HOI4D_overall/ZY20210800002/H2/C1/N01/S336/s02/T1",
    # ]
    # frame_ids = np.arange(0, 300)
    # category = "rigid/ToyCar"
    # sym_axis = None
    # use_gt_model = True  # use gt model or completed model
    # # -----------------------------------------------------------------------------------------------------------
    # # -----------------------------------------------------------------------------------------------------------
    # # chair
    # video_dirs = [
    #     "/share/datasets/HOI4D_overall/ZY20210800001/H1/C20/N18/S286/s01/T1",
    #     "/share/datasets/HOI4D_overall/ZY20210800001/H1/C20/N20/S381/s03/T3",
    #     "/share/datasets/HOI4D_overall/ZY20210800001/H1/C20/N25/S289/s02/T2",
    #     "/share/datasets/HOI4D_overall/ZY20210800001/H1/C20/N30/S295/s04/T4",
    #     "/share/datasets/HOI4D_overall/ZY20210800001/H1/C20/N47/S381/s03/T3",
    # ]
    # frame_ids = np.arange(0, 300)
    # category = "rigid/Chair"
    # sym_axis = None
    # use_gt_model = True  # use gt model or completed model
    # # # -----------------------------------------------------------------------------------------------------------
    # # mug
    # video_dirs = [
    #     "/share/datasets/HOI4D_overall/ZY20210800001/H1/C2/N17/S212/s05/T3",
    #     "/share/datasets/HOI4D_overall/ZY20210800001/H1/C2/N25/S213/s02/T4",
    #     "/share/datasets/HOI4D_overall/ZY20210800001/H1/C2/N30/S92/s04/T1",
    #     "/share/datasets/HOI4D_overall/ZY20210800002/H2/C2/N39/S30/s05/T1",
    #     "/share/datasets/HOI4D_overall/ZY20210800003/H3/C2/N42/S217/s05/T2",
    # ]
    # frame_ids = np.arange(0, 300)
    # category = "rigid/Mug"
    # sym_axis = None
    # use_gt_model = True  # use gt model or completed model
    # # -----------------------------------------------------------------------------------------------------------
    # # -----------------------------------------------------------------------------------------------------------
    # # bottle
    # video_dirs = [
    #     "/share/datasets/HOI4D_overall/ZY20210800001/H1/C5/N46/S57/s05/T1",
    #     "/share/datasets/HOI4D_overall/ZY20210800002/H2/C5/N14/S255/s04/T1",
    #     "/share/datasets/HOI4D_overall/ZY20210800002/H2/C5/N17/S56/s01/T2",
    #     "/share/datasets/HOI4D_overall/ZY20210800002/H2/C5/N24/S268/s01/T6",
    #     "/share/datasets/HOI4D_overall/ZY20210800002/H2/C5/N40/S265/s03/T4",
    # ]
    # frame_ids = np.arange(0, 300)
    # category = "rigid/Bottle"
    # sym_axis = "z"
    # use_gt_model = True  # use gt model or completed model
    # # -----------------------------------------------------------------------------------------------------------
    # # -----------------------------------------------------------------------------------------------------------
    # # bowl
    # video_dirs = [
    #     "/share/datasets/HOI4D_overall/ZY20210800001/H1/C7/N14/S280/s04/T5",
    #     "/share/datasets/HOI4D_overall/ZY20210800001/H1/C7/N17/S276/s02/T3",
    #     "/share/datasets/HOI4D_overall/ZY20210800001/H1/C7/N20/S59/s02/T1",
    #     "/share/datasets/HOI4D_overall/ZY20210800001/H1/C7/N25/S60/s01/T2",
    #     "/share/datasets/HOI4D_overall/ZY20210800004/H4/C7/N13/S57/s01/T1",
    # ]
    # frame_ids = np.arange(0, 300)
    # category = "rigid/Bowl"
    # sym_axis = "z"
    # use_gt_model = True  # use gt model or completed model
    # # -----------------------------------------------------------------------------------------------------------
    
    
    r_errs_all, t_errs_all = None, None
    for video_dir in video_dirs:
        # gt_model_path = "/share/datasets/HOI4D_CAD_Model_for_release/rigid/Kettle/031.obj"
        p0 = video_dir.find("N")
        p1 = video_dir.find("/", p0)
        N_id = int(video_dir[p0+1:p1])
        gt_model_path = os.path.join("/share/datasets/HOI4D_CAD_Model_for_release", category, str(N_id).zfill(3) + ".obj")
        completed_model_path = "/home/liuyun/HOI-Sim/object_pose_tracking/pcd_completion/HOI4D_kettle/train_visualization/100.ply"  # TODO
        if use_gt_model:
            model_points = read_model_pcd(gt_model_path, N_points=2000)
        else:
            init_pose_path = os.path.join(video_dir, "objpose", "0.json")
            if not os.path.isfile(init_pose_path):
                init_pose_path = os.path.join(video_dir, "objpose", "00000.json")
            init_pose = read_pose_from_json(init_pose_path)
            model_pcd = o3d.io.read_point_cloud(completed_model_path)
            model_points = np.float32(model_pcd.points) * init_pose["scale"].reshape(3)
        r_errs, t_errs = track_ShapeAlign(video_dir, frame_ids, model_points, sym_axis=sym_axis)
        if not r_errs_all is None:
            r_errs_all = np.concatenate((r_errs_all, r_errs))
            t_errs_all = np.concatenate((t_errs_all, t_errs))
        else:
            r_errs_all = r_errs
            t_errs_all = t_errs
        
    statistics_path = os.path.join("./results", "statistics_" + video_dirs[0].split("/")[-5] + ".txt")
    wr = open(statistics_path, "w")
    wr.write("r_err(deg) = " + str(np.mean(r_errs_all)) + " (+-" + str(np.std(r_errs_all)) +  ")\n")
    wr.write("t_err(cm) = " + str(np.mean(t_errs_all)) + " (+-" + str(np.std(t_errs_all)) +  ")\n")
    wr.close()
